from instruments.instrument import Instrument

from utils import frequency_unit_conversion



class KeysightSignalGenerator(Instrument):
    def __init__(self, address, parent=None):
        super().__init__(address, parent)

    def set_frequency(self, frequency_hz):
        self._inst.write(":FREQuency {0}".format(frequency_hz))

    def get_frequency(self, target_unit: str = 'MHz'):
        frequency = self._inst.query_cmd(":FREQuency?")
        frequency = frequency_unit_conversion(frequency, "Hz", target_unit)
        return frequency

    def set_amplitude(self, amplitude):
        self._inst.write(":POWer {0} dBm".format(amplitude))

    def get_power(self):
        scpi = ':POWer?'
        power = self._inst.query_cmd(scpi)
        return round(float(power), 2)

    def set_rf(self, is_rf_on: int):
        stat = ['OFF', 'ON']
        self._inst.write("OUTP:STAT " + stat[is_rf_on])


class RSSignalGenerator(Instrument):
    def __init__(self, address, parent=None):
        super().__init__(address, parent)

    def set_frequency(self, frequency_mhz):
        # self._inst.write_cmd(":SOURce:FREQuency:CW {0}".format(frequency_mhz * 1000000))
        self._inst.write_cmd("FREQUENCY {0}".format(frequency_mhz))

    def get_frequency(self, target_unit: str = 'MHz'):
        frequency = self._inst.query_cmd("FREQUENCY?")
        frequency = frequency_unit_conversion(frequency, "Hz", target_unit)
        return frequency

    def set_amplitude(self, amplitude):
        # self._inst.write_cmd(":SOURce:POWer:LEVel:IMMediate:AMPLitude {0} dBm".format(amplitude))
        self._inst.write_cmd("POWER {0}".format(amplitude))

    def get_power(self):
        scpi = ":SOURce1:POWer:LEVel:IMMediate:AMPLitude?"
        power = self._inst.query_cmd(scpi)
        return round(float(power), 2)

    def set_rf(self, is_rf_on: int):
        if is_rf_on == 0:
            state = 'OFF'
        else:
            state = 'ON'
        self._inst.write_cmd("OUTPUT {0}".format(state))